#include "protoserv/compiler/rpc_plugin.h"

#ifdef _WIN32
#  include <io.h>
#  include <fcntl.h>
#  ifndef STDIN_FILENO
#    define STDIN_FILENO 0
#  endif
#  ifndef STDOUT_FILENO
#    define STDOUT_FILENO 1
#  endif
#else
#  include <unistd.h>
#endif

#include <stdio.h>

#include <iostream>
#include <limits>
#include <set>
#include <vector>

#include <google/protobuf/compiler/code_generator.h>
#include <google/protobuf/compiler/plugin.pb.h>
#include <google/protobuf/descriptor.h>
#include <google/protobuf/descriptor.pb.h>
#include <google/protobuf/stubs/strutil.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <google/protobuf/stubs/common.h>
#include "protoserv/protocol.pb.h"

namespace protoserv {
namespace compiler {

using namespace google::protobuf::compiler;
using namespace google::protobuf;
using namespace std;

static const int kMinReservedExtension = 19000;
static const int kMaxReservedExtension = 19999;
static const int kMinExtension = 16;
static const int kMaxExtension = 536870911;

class GeneratorResponseContext : public GeneratorContext {
 public:
  GeneratorResponseContext(CodeGeneratorResponse* response,
                           const vector<const FileDescriptor*>& parsed_files)
      : response_(response),
        parsed_files_(parsed_files) {}
  virtual ~GeneratorResponseContext() {}

  // implements GeneratorContext --------------------------------------

  virtual io::ZeroCopyOutputStream* Open(const string& filename) {
    CodeGeneratorResponse::File* file = response_->add_file();
    file->set_name(filename);
    return new io::StringOutputStream(file->mutable_content());
  }

  virtual io::ZeroCopyOutputStream* OpenForInsert(
      const string& filename, const string& insertion_point) {
    CodeGeneratorResponse::File* file = response_->add_file();
    file->set_name(filename);
    file->set_insertion_point(insertion_point);
    return new io::StringOutputStream(file->mutable_content());
  }

  void ListParsedFiles(vector<const FileDescriptor*>* output) {
    *output = parsed_files_;
  }

 private:
  CodeGeneratorResponse* response_;
  const vector<const FileDescriptor*>& parsed_files_;
};

unsigned int Adler32(const string& text) {
  static const int kModAdler = 65521;
  const char* data = text.data();
  unsigned int a = 1, b = 0;
  for (size_t i = 0; i < text.size(); ++i) {
    a = (a + data[i]) % kModAdler;
    b = (b + a) % kModAdler;
  }
  return (b << 16) | a;
}

bool IsMethodIdValid(int method_id) {
  if (method_id < kMinExtension || method_id >= kMaxExtension) {
    return false;
  }
  if (method_id >= kMinReservedExtension &&
      method_id <= kMaxReservedExtension) {
    return false;
  }
  return true;
}

int NormalizeToExtensionIdRange(unsigned int hash) {
  // Number reserved extension ids.
  static const int kTotalReservedExtensions =
      kMaxReservedExtension - kMinReservedExtension + 1;
  // Number of valid extension ids.
  static const double kTotalValidExtensionIds =
      (kMaxExtension - kMinExtension + 1) - kTotalReservedExtensions;
  // Brings hash to extension id domain size.
  const double raw_id =
      kTotalValidExtensionIds * (hash / static_cast<double>(UINT_MAX));
  // Adds first offset.
  int id = static_cast<int>(raw_id) + kMinExtension;
  if (id >= kMinReservedExtension) {
    // Skips reserved range.
    id += kTotalReservedExtensions;
  }
  if (!IsMethodIdValid(id)) {
    fprintf(stderr, "Invalid id: %u, hash=%u\n", id, hash);
    assert(false);
  }
  return id;
}

int ComputeMethodId(const string& method_name) {
  return NormalizeToExtensionIdRange(Adler32(method_name));
}

int ComputeMethodId(const MethodDescriptor* method) {
  string method_name;
  method_name.append(method->service()->full_name());
  method_name.append(".");
  method_name.append(method->name());
  return ComputeMethodId(method_name);
}

void InitRpcExtension(int number,
                      const string& extendee,
                      FieldDescriptorProto* extension) {
  extension->set_number(number);
  extension->set_extendee(extendee);
  extension->set_label(FieldDescriptorProto::LABEL_OPTIONAL);
  extension->set_type(FieldDescriptorProto::TYPE_MESSAGE);
}

void AddMethodRpcExtensions(const MethodDescriptor* method,
                            FileDescriptorProto* proto) {
  int number = method->options().GetExtension(protoserv::method_id);
  if (number == 0) {
    number = ComputeMethodId(method);
  } else if (!IsMethodIdValid(number)) {
    fprintf(stderr, "Invalid method id: %d", number);
    exit(1);
  }
  const string extension_prefix =
      method->service()->name() + "_" + method->name();
  FieldDescriptorProto* request_ext = proto->add_extension();
  InitRpcExtension(number, "protoserv.RequestProto", request_ext);
  request_ext->set_name(extension_prefix + "_Input");
  request_ext->set_type_name(method->input_type()->full_name());

  FieldDescriptorProto* response_ext = proto->add_extension();
  InitRpcExtension(number, "protoserv.ResponseProto", response_ext);
  response_ext->set_name(extension_prefix + "_Output");
  response_ext->set_type_name(method->output_type()->full_name());
}

string StripProto(const string& filename) {
  if (HasSuffixString(filename, ".protodevel")) {
    return StripSuffixString(filename, ".protodevel");
  } else {
    return StripSuffixString(filename, ".proto");
  }
}


void FillRpcExtenionsFile(const FileDescriptor* file,
                          FileDescriptorProto* proto) {
  FileDescriptorProto orig_proto;
  file->CopyTo(&orig_proto);
  proto->set_name(StripProto(orig_proto.name()) + "_rpc.proto");
  proto->set_package(orig_proto.package());
  proto->mutable_options()->CopyFrom(orig_proto.options());
  proto->add_dependency("protoserv/protocol.proto");
  proto->add_dependency(orig_proto.name());
  for (int i = 0; i < file->service_count(); ++i) {
    const ServiceDescriptor* service = file->service(i);
    for (int j = 0; j < service->method_count(); ++j) {
      AddMethodRpcExtensions(service->method(j), proto);
    }
  }
}

void GenerateFiles(const vector<const FileDescriptor*>& files,
                   const CodeGeneratorRequest& request,
                   CodeGenerator* generator,
                   GeneratorResponseContext* context,
                   CodeGeneratorResponse* response) {
  for (size_t i = 0; i < files.size(); i++) {
    const FileDescriptor* file = files[i];
    string error;
    bool succeeded = generator->Generate(
        file, request.parameter(), context, &error);
    if (!succeeded && error.empty()) {
      error = "Code generator returned false but provided no error "
              "description.";
    }
    if (!error.empty()) {
      response->set_error(file->name() + ": " + error);
      break;
    }
  }
}

int ProtoservPlugin::Main(
    int argc, char* argv[],
    google::protobuf::compiler::CodeGenerator* ext_generator,
    google::protobuf::compiler::CodeGenerator* rpc_generator) {
  if (argc > 1) {
    cerr << argv[0] << ": Unknown option: " << argv[1] << endl;
    return 1;
  }

#ifdef _WIN32
  _setmode(STDIN_FILENO, _O_BINARY);
  _setmode(STDOUT_FILENO, _O_BINARY);
#endif

  CodeGeneratorRequest request;
  if (!request.ParseFromFileDescriptor(STDIN_FILENO)) {
    cerr << argv[0] << ": protoc sent unparseable request to plugin." << endl;
    return 1;
  }

  DescriptorPool pool;
  for (int i = 0; i < request.proto_file_size(); i++) {
    const FileDescriptor* file = pool.BuildFile(request.proto_file(i));
    if (file == NULL) {
      // BuildFile() already wrote an error message.
      return 1;
    }
  }

  vector<const FileDescriptor*> parsed_files;
  vector<const FileDescriptor*> ext_files;
  for (int i = 0; i < request.file_to_generate_size(); i++) {
    const FileDescriptor* file =
        pool.FindFileByName(request.file_to_generate(i));
    if (file == NULL) {
      cerr << argv[0] << ": protoc asked plugin to generate a file but "
              "did not provide a descriptor for the file: "
           << request.file_to_generate(i) << endl;
      return 1;
    }
    parsed_files.push_back(file);
    if (file->service_count() > 0) {
      FileDescriptorProto ext_file_proto;
      FillRpcExtenionsFile(file, &ext_file_proto);
      const FileDescriptor* ext_file = pool.BuildFile(ext_file_proto);
      if (ext_file == NULL) {
        // BuildFile() already wrote an error message.
        return 1;
      }
      parsed_files.push_back(ext_file);
      ext_files.push_back(ext_file);
    }
  }

  CodeGeneratorResponse response;
  GeneratorResponseContext context(&response, parsed_files);
  GenerateFiles(parsed_files, request, rpc_generator, &context, &response);
  GenerateFiles(ext_files, request, ext_generator, &context, &response);

  if (!response.SerializeToFileDescriptor(STDOUT_FILENO)) {
    cerr << argv[0] << ": Error writing to stdout." << endl;
    return 1;
  }

  return 0;
}

}  // namespace compiler
}  // namespace protoserv
